#!/bin/bash
function showhelp() {
    echo 'usage: cp2kCDD abc.pdb will do a simple structure analysis, give the z coord range for each element'
    echo 'cp2kCDD.sh -z 0.5 pdb/cif/restart will divided the input structure'  
    echo '-z is z distance in frac_coord, above which is fragment A and below which is fragment B'
    echo 'run this script in a completed cp2k job dir with a optimized pdb file, it will generate FRAG_A, FRAG_B and FRAG_AB dir'
    echo 'run this script again with finished job in FRAG_A, FRAG_B and FRAG_AB dir will output CDD, workfunction, dipole, charge, IGMH'
    echo '-c 0 1 -1 2 1 2 set charge and multiplicity for AB A B, respectively'
    echo '-R regenerate input file even if there are fininshed job'
    echo '-G set IGMH grid quality, 1:low, 2:medium, 3:high, 0:no igmh (default)'
    echo '-n set number of threads used by Multiwfn, default all threads'
    echo '-C set atomic charge type 1:Mulliken, 2:Lowdin, 3:Hirshfeld (default), 4:Hirshfeld-i, 5:voronoi, 6:resp, 7:REPEAT '
    exit 1
}
# todo list
# read keywords from restart file
#  
if [ $# -lt 1 ]; then
    showhelp
fi
while getopts ":z:t:m:b:s:c:S:k:C:G:n:RuT" opt
do
    case $opt in
        z) zth=$OPTARG;;
        R) RECALC=on;;
        G) GRID=$OPTARG;;
        n) NTHREAD=$OPTARG;;
        m) METHOD=$OPTARG;;
        b) BASIS=$OPTARG;;
        s) SCF=$OPTARG;;
        c) CHARGEMULT=$OPTARG;;
        C) CHARGETYPE=$OPTARG;;
        S) SCCS=$OPTARG;;
        k) KPOINT=$OPTARG;;
        u) DFTU=on;;
        ?) showhelp;;
    esac 
done
shift $(($OPTIND - 1))

function gen_minp {
    # $1 is input name $2 is output name $3 is charge $4 is multiplicity
    echo "cp2k" > MULTIWFN_INPUT
    echo "$2" >> MULTIWFN_INPUT
    echo -en "-9\n1\n$3\n" >> MULTIWFN_INPUT
    echo -en "2\n$4\n" >> MULTIWFN_INPUT
    if [[ "$DFTU" == "on" ]];then
      echo 8 >> MULTIWFN_INPUT
    fi
    echo "0" >> MULTIWFN_INPUT
    # Set charge type to calculate Hirshfeld
    echo "-4" >> MULTIWFN_INPUT
    echo "$CHARGETYPE" >> MULTIWFN_INPUT
    # Set exporting cube file: den and pot
    echo -en "-3\n7\n" >> MULTIWFN_INPUT
    # default exporting molden file
    echo "-2" >> MULTIWFN_INPUT
    # choose task
    echo -en "-1\n1\n" >> MULTIWFN_INPUT
    # choose theoretical method
    echo -en "1\n$METHOD\n" >> MULTIWFN_INPUT
    # choose basis set
    echo -en "2\n$BASIS\n" >> MULTIWFN_INPUT
    # choose dispersion correction
    echo -en "3\n2\n" >> MULTIWFN_INPUT
    # chooe SCF strategy
    if [[ -n "$SCF" ]];then
        echo -en "$SCF\n" >> MULTIWFN_INPUT
    fi
    # SCCS
    if [[ -n "$SCCS" ]];then
      echo "7" >> MULTIWFN_INPUT
    fi
    # turn on smear
      echo "6" >> MULTIWFN_INPUT
    # KPOINT 
    echo -en "8\n$KPOINT\n" >> MULTIWFN_INPUT
    # outputted structure format
    echo -en "-5\n3\n" >> MULTIWFN_INPUT
    if [[ -n "$FIXATOM" ]];then
       echo -en "9\n$FIXATOM\n" >> MULTIWFN_INPUT
    fi  
    echo -en "0\nq" >> MULTIWFN_INPUT
    cat MULTIWFN_INPUT | Multiwfn $1 1> /dev/null
    sed -i '/steps of inner SCF/c\      MAX_SCF 90' $2
    sed -i '/steps of outer SCF/c\        MAX_SCF 4' $2
    sed -i '/ PRINT_LEVEL/c\  PRINT_LEVEL MEDIUM' $2
}

function wrap_fcoord {
    # wrap fraction z coord of POSCAR.cif generated by Multiwfn
    awk '{if(NR > 16){
              if($4 > 1) { 
                  $4 = $4 - int($4) 
                  $0 = sprintf("%2s %6.8f %6.8f %6.8f",$1,$2,$3,$4)
              } 
              else if ($4 < 0) {
                  $4 = 1 - int($4) + $4
                  $0 = sprintf("%2s %6.8f %6.8f %6.8f",$1,$2,$3,$4)
              }
          }
          print $0 }' POSCAR.cif > cp2kCDD_tmp.cif
    mv cp2kCDD_tmp.cif POSCAR.cif
}

function init_multiwfn {
    mpath=$(dirname `which Multiwfn`)
    if [ -f $mpath/settings.ini ];then
        cp $mpath/settings.ini ./
    else
        echo "Warning! settings.ini not found in $mpath, I could not modify settings"
        return 1
    fi
    sysproc=`awk -F: '/^Core\(s\) per socket:/ {core=$2} /^Socket\(s\):/ {socket=$2} END {print core*socket}' < <(lscpu)`
    if [ -z $NTHREADS ];then
        NTHREADS=$sysproc
    fi
    echo "set Multiwfn nthreads to $NTHREADS"
    sed -i "/nthreads=/c\nthreads= $NTHREADS" settings.ini
}

function patch_molden {
    # add [cell] and [Nval] section to the head of molden and generate a patched_molden file
    # $1 is the cp2k input file and  $2 is the molden_file need to be patched
    awk 'BEGIN{print "[Cell]"}{
        if($0 ~/&CELL/)
            rf=1;
        if(rf==1){
            if(NF==4){
                printf "%12.6f%12.6f%12.6f\n",$2,$3,$4
            }
        }
        if($0~/&END CELL/){
            rf=0
        }}' $1 > molden_patch.txt

    awk 'BEGIN{print "[Nval]"}{
        if($0 ~/&KIND/)
            rf=1;
        if(rf==1){
            if($0 ~ /BASIS_SET.*-q([0-9]+)/){
                match($0, /-q([0-9]+)/, arr)
            }
            if($0 ~ /ELEMENT/){
                ele=$2
            }
        }
        if($0~/&END KIND/){
            printf "%-3s%3s\n",ele,arr[1]
            rf=0
        }}END{print ""}' $1 >> molden_patch.txt
    cat molden_patch.txt $2 > ${2/.molden/}_patched.molden 
}



function work_func {
    # comput fermi level and  work_function 
    # $1 is basename 
    fermi=`awk '/Fermi energy/{print $NF*27.211386245988}' ${1}.out |tail -n 1`
    echo -en "13\n18\nZ\na\n-1\n6\n9\n0\n-1\nq" |  Multiwfn ${1}-v_hartree-1_0.cube &> /dev/null
    vacuum=`awk '{idx+=1;e[idx]=$3}END{
             delta=100
             for(i=2;i<length(e);i++){
                 d1=e[i]-e[i-1]
                 d2=e[i]-e[i+1]
                 sub(/^-/,"",d1);sub(/^-/,"",d2)
                 if(d1+d2<delta) {delta=d1+d2;min_idx=i}
             }print e[min_idx]*27.211386245988
         }' pleavgcurve.txt`
    work_func=`echo "$vacuum - $fermi" | bc`
}

function calc_igmh {
    # use Multiwfn to calculate IGMH between surface and absorpant
    if [[ $GRID == 0 ]];then
        return 0
    fi
    if [[ ! -f "${basename}-inter_igmh.cub" || ! -f "${basename}_sl2r.cub" ]];then
        idxA=`cat FRAG_A_INDEX.ndx`
        idxB=`cat FRAG_B_INDEX.ndx`
        if [ -z "$idxA" ];then
            echo "Error! Index of FRAG A is empty, check FRAG_A_INDEX.ndx file"
            return 1
        fi
        if [ -z "$idxB" ];then
            echo "Error! Index of FRAG B is empty, check FRAG_B_INDEX.ndx file"
            return 1
        fi
        if [ ! -e "FRAG_AB/AB-MOS-1_0.molden" ];then
            echo "Error! AB-MOS-1_0.molden not found in FRAG_AB"
            return 1
        fi
        echo "Starting generating igmh cube ..."
        # generate Multiwfn command
        cd FRAG_AB
        patch_molden AB.inp AB-MOS-1_0.molden
        cd ..
        echo -en "\n\n20\n11\n2\n${idxA}\nc\n${GRID}\n3\n2\n1\n1\n0\n0\nq" | Multiwfn FRAG_AB/AB-MOS-1_0_patched.molden |  awk -v RS='[\n\r]'  '{if($0 ~/Progress:/){pflag=1;printf "%s\r",$0}if(pflag==    1 && $0 !~/Progress:/){print "";pflag=0}}'
        mv dislin.png ${basename}_igmh.png
        mv dg_inter.cub ${basename}-inter_igmh.cub
        echo q | Multiwfn FRAG_AB/AB.inp  | grep "pbc set {" > ${basename}-inter_igmh.pbc
        mv output.txt ${basename}-igmh.txt
        mv dg_intra.cub ${basename}-intra_igmh.cub
        echo q | Multiwfn FRAG_AB/AB.inp  | grep "pbc set {" > ${basename}-intra_igmh.pbc
        mv sl2r.cub ${basename}_sl2r.cub
         
    else
        echo q | Multiwfn FRAG_AB/AB.inp  | grep "pbc set {" > ${basename}-inter_igmh.pbc
        echo "${basename}-inter_igmh.cub and ${basename}_sl2r.cub already exists, skip the calculation, you need to remove this file to regenerate it"
    fi
}

if [[ -z "$METHOD" ]];then METHOD=2;fi
if [[ -z "$BASIS" ]];then BASIS=2;fi
if [[ -z "$GRID" ]];then GRID=3;fi
#if [[ -z "$SCF" ]];then
#    if [[ "$METHOD" == "30" ]];then
#    SCF=5
#    else
#    SCF="4"
#    fi
#fi
if [[ -z "$CHARGEMULT" ]];then CHARGEMULT="0 1 0 1 0 1";fi
if [[ -z "$KPOINT" ]];then KPOINT=1,1,1;fi
if [[ -z "$CHARGETYPE" ]];then CHARGETYPE=3;fi
CHARGEAB=`echo $CHARGEMULT | cut -d " " -f 1`
MULTIAB=`echo $CHARGEMULT | cut -d " " -f 2`
CHARGEA=`echo $CHARGEMULT | cut -d " " -f 3`
MULTIA=`echo $CHARGEMULT | cut -d " " -f 4`
if [[ -z $CHARGEA ]];then CHARGEB=$CHARGEAB;fi
if [[ -z $MULTIA ]];then MULTIA=$MULTIAB;fi
CHARGEAB=`echo $CHARGEMULT | cut -d " " -f 5`
MULTIAB=`echo $CHARGEMULT | cut -d " " -f 6`
if [[ -z $CHARGEB ]];then CHARGEB=$(($CHARGEAB-$CHARGEA));fi
if [[ -z $MULTIB ]];then MULTIB=$(($MULTIAB-1-$MULTIA+1+1));fi


echo "system surf_dipole_ave surf_dipole_sum binding_ene charge_A charge_B work_func_A work_func_B work_func_AB" > ALLCDDRESULT.txt
wd=`pwd`
while [[ $# -ge 1 ]]
do
if [ -d "$1" ];then
    ISDIR=1
    cd $1
    INP=`ls *.restart | tail -1`
    basename=${INP%.*}
else
    ISDIR=0
    INP=$1
    basename=${INP%.*}
fi
# check is job is completed 
if [[ -f "FRAG_A/A-ELECTRON_DENSITY-1_0.cube" ]] && [[ -f "FRAG_B/B-ELECTRON_DENSITY-1_0.cube" ]] && [[ -f "FRAG_AB/AB-ELECTRON_DENSITY-1_0.cube" ]] && [[ $RECALC != "on" ]] && [ -f ZTH_FOR_FRAG ]; then
    zth=`cat ZTH_FOR_FRAG`
    init_multiwfn
    echo "found zth for separate A and B is $zth"
    # generate CDD cube
    if [ ! -f "${basename}-CDD.cub" ];then
        echo -en "13\n11\n4\nFRAG_A/A-ELECTRON_DENSITY-1_0.cube\n11\n4\nFRAG_B/B-ELECTRON_DENSITY-1_0.cube\n" > MULTIWFN_COM 
        echo -en "0\nCDD.cub\n18\nZ\na\n-1\n5\n8\n9\n0\n-1\nq" >> MULTIWFN_COM 
        cat MULTIWFN_COM | Multiwfn FRAG_AB/AB-ELECTRON_DENSITY-1_0.cube | awk -v RS='[\n\r]'  '{if($0 ~/Progress:/){pflag=1;printf "%s\r",$0}if(pflag==1 && $0 !~/Progress:/){print "";pflag=0}}'
        mv dislin.png ${basename}.png
        mv CDD.cub ${basename}-CDD.cub
        mv locintcurve.txt ${basename}-CDD.txt
        echo q | Multiwfn FRAG_AB/AB.inp  | grep "pbc set {" > ${basename}-CDD.pbc
    else
        echo "${basename}-CDD.cub already exists, skip the calculation, you need to remove this file to regenerate it"
    fi
    # comput dipole moment 
    area=`echo q | Multiwfn $INP | awk '/pbc set/{area=$4*$5*sin($9/180*atan2(0,-1));print area}'`
    dpa=`awk '{d+=$2*$3}END{print d*4.8032}' pleavgcurve.txt`
    dpi=`awk -v A=$area '{d+=$2*$3}END{print d*4.8032/A}' ${basename}-CDD.txt`
    echo "${basename}: dipole from plane averaged CDD/plane summed CDD: $dpa $dpi"
    echo -n "${basename} $dpa $dpi" >> "${wd}/ALLCDDRESULT.txt"
    # compute binding energy charge transfer
    eneA=`grep 'Total energy:' FRAG_A/A.out | tail -n 1 | awk '{print $3}'`
    eneA=`echo "$eneA * 27.211396" | bc`
    eneB=`grep 'Total energy:' FRAG_B/B.out | tail -n 1 | awk '{print $3}'`
    eneB=`echo "$eneB * 27.211396" | bc`
    eneAB=`grep 'Total energy:' FRAG_AB/AB.out | tail -n 1 | awk '{print $3}'`
    eneAB=`echo "$eneAB * 27.211396" | bc`
    be=`echo "$eneAB - $eneA - $eneB" | bc`
    printf "${basename} Binding Energy(BE) = AB - A - B is:\n"
    printf "BE = %.6f - %.6f - %.6f = %.6f eV\n" $eneAB $eneA $eneB $be
    echo -n " $be" >> "${wd}/ALLCDDRESULT.txt"
    # compute hirshfeld charge transfer
    rm HIRSHFELD_* 2>/dev/null
    awk 'BEGIN{hc_flag=0;fidx=0;nr=-100}{if($0 ~/Hirshfeld Charges/){nr=NR;fidx=1};if(NR==nr+3){hc_flag=1};if(hc_flag==1){if(NF<6) hc_flag=0; else {hc[fidx] += $NF;fidx+=1}}}END{for (c in hc) print hc[c] > "HIRSHFELD_CHARGE"}' FRAG_AB/AB.out
    paste COORD_AB.txt HIRSHFELD_CHARGE > HIRSHFELD_AB.txt
    zth=`cat ZTH_FOR_FRAG`
    awk -v z=$zth 'BEGIN{upQ=0;downQ=0}
    {if(NF==5){
        if($4 > z){
            upQ += $5
            countA[$1]++
        } else {
            downQ += $5
            countB[$1]++
        }
    }}
    END{
        for (c in countA) fa=fa""sprintf("%s%s",c,countA[c])
        for (c in countB) fb=fb""sprintf("%s%s",c,countB[c])
        printf "FRAG_A (%20s) Hirshfeld charge: %8.3f\nFrag_B (%20s) Hirshfeld charge: %8.3f\n",fa,upQ,fb,downQ}' HIRSHFELD_AB.txt | tee CHARGE_TEMP
    echo -n " "$(awk '{print $NF}' CHARGE_TEMP | tr '\n' ' ') >> "${wd}/ALLCDDRESULT.txt"
    # computer fermi and work function
    fermiA=`awk '/Fermi Energy/{print $NF}' FRAG_A/A.out |tail -n 1`
    fermiB=`awk '/Fermi Energy/{print $NF}' FRAG_B/B.out |tail -n 1`
    fermiAB=`awk '/Fermi Energy/{print $NF}' FRAG_AB/AB.out |tail -n 1`
    printf "%-10s%10s%10s%10s\n" frag fermi vacuum work_func
    cd FRAG_A
    work_func A
    printf "%-10s%10.3f%10.3f%10.3f\n" FRAG_A $fermi $vacuum $work_func
    echo -n " $work_func" >> "${wd}/ALLCDDRESULT.txt"
    cd ../FRAG_B
    work_func B
    printf "%-10s%10.3f%10.3f%10.3f\n" FRAG_B $fermi $vacuum $work_func
    echo -n " $work_func" >> "${wd}/ALLCDDRESULT.txt"
    cd ../FRAG_AB
    work_func AB
    printf "%-10s%10.3f%10.3f%10.3f\n" FRAG_AB $fermi $vacuum $work_func
    echo -n " $work_func" >> "${wd}/ALLCDDRESULT.txt"
    cd ..
    # comput igmh
    echo  "" >> "${wd}/ALLCDDRESULT.txt"
    calc_igmh
else
#if [[ -z "$zth" ]];then
#     echo -en "100\n2\n33\nPOSCAR.gro\n0\nq" |  Multiwfn $INP &> /dev/null
#     awk '{if(NR>2 && NF < 7){
#               maxarr[$2]=($6 > maxarr[$2])?$6:maxarr[$2]
#               if($2 in minarr) {
#                   minarr[$2]=($6 < minarr[$2])?$6:minarr[$2]
#               } else {
#                   minarr[$2]=$6
#               }
#           }}END{
#               for(i in minarr) printf "%5s%10.6f%10.6f\n",i,minarr[i],maxarr[i]}' POSCAR.gro
#     exit
#     rm POSCAR.gro
#fi

    # print structure info
    if [[ -z "$zth" ]];then
         echo -en "100\n2\n33\nPOSCAR.cif\n0\nq" |  Multiwfn $INP &> /dev/null
         wrap_fcoord
         awk 'BEGIN{printf "%15s%10s\n","min","max"}{if(NR>16 && NF ==4){
                   maxarr[$1]=($4 > maxarr[$1])?$4:maxarr[$1]
                   if($1 in minarr) {
                       minarr[$1]=($4 < minarr[$1])?$4:minarr[$1]
                   } else {
                       minarr[$1]=$4
                   }
               }}END{
                   for(i in minarr) printf "%5s%10.6f%10.6f\n",i,minarr[i],maxarr[i]}' POSCAR.cif
         exit
         cat POSCAR.cif
    fi
    mkdir FRAG_A
    mkdir FRAG_B
    mkdir FRAG_AB
#echo -en "100\n2\n34\nPOSCAR.gro\n0\nq" |  Multiwfn $INP &> /dev/null
#awk -v z=$zth 'BEGIN{idxa=1;idxb=1;ncoordsa=0;ncoordsb=0}
#     {if(NF==6 && NR > 2){
#        if($6 > z){
#            linesA[idxa]=$0;idxa++;ncoordsa++
#        } else {
#            linesB[idxb]=$0;idxb++;ncoordsb++
#        }
#     }else{
#        linesA[idxa]=$0;idxa++
#        linesB[idxb]=$0;idxb++
#     }}
#     END{
#        linesA[2]=ncoordsa
#        linesB[2]=ncoordsb
#        printf "%s","" > "FRAG_A.gro"
#        printf "%s","" > "FRAG_B.gro"
#        for(i=1;i<=idxa+1;i++){
#            print linesA[i] >> "FRAG_A.gro"}
#        for(i=1;i<=idxb+1;i++){
#            print linesB[i] >> "FRAG_B.gro"}
#        }' POSCAR.gro
#FA=`awk '{ if(NF==6 && NR > 2) count[$2]++ } END { for (c in count) printf "%s%s",c,count[c] }' FRAG_A.gro`
#FB=`awk '{ if(NF==6 && NR > 2) count[$2]++ } END { for (c in count) printf "%s%s",c,count[c] }' FRAG_B.gro`
#mv FRAG_A.gro FRAG_A
#mv FRAG_B.gro FRAG_B
#mv POSCAR.gro FRAG_AB
    echo -en "100\n2\n33\nPOSCAR.cif\n0\nq" |  Multiwfn $INP &> /dev/null
    wrap_fcoord
    awk -v z=$zth 'BEGIN{idxa=1;idxb=1;idxab=1;idxA="";idxB="";ncoordsa=0;ncoordsb=0}
         {if(NF==4 && NR > 16){
            if($4 > z){
                linesA[idxa]=$0;idxa++;ncoordsa++
                if (idxA == "")
                    idxA = idxab
                else
                    idxA=idxA "," idxab
            } else {
                linesB[idxb]=$0;idxb++;ncoordsb++
                if (idxB == "")
                    idxB = idxab
                else
                    idxB=idxB "," idxab
            }
            coordsAB[idxab]=$0;idxab++
         }else{
            linesA[idxa]=$0;idxa++
            linesB[idxb]=$0;idxb++
         }}
         END{
            printf "%s","" > "FRAG_A.cif"
            printf "%s","" > "FRAG_B.cif"
            printf "%s","" > "COORD_AB.txt"
            printf "%s\n",idxA > "FRAG_A_INDEX.ndx"
            printf "%s\n",idxB > "FRAG_B_INDEX.ndx"
            for(i=1;i<=idxa+1;i++){
                print linesA[i] >> "FRAG_A.cif"}
            for(i=1;i<=idxb+1;i++){
                print linesB[i] >> "FRAG_B.cif"}
            for(i=1;i<=idxab+1;i++){
                print coordsAB[i] >> "COORD_AB.txt"}
            }' POSCAR.cif
    mv FRAG_A.cif FRAG_A
    mv FRAG_B.cif FRAG_B
    mv POSCAR.cif FRAG_AB
    echo $zth > ZTH_FOR_FRAG

    cd FRAG_A
    gen_minp FRAG_A.cif A.inp $CHARGEA $MULTIA
    FA=`awk '{ if(NF==4 && NR > 16) count[$1]++ } END { for (c in count) printf "%s%s",c,count[c] }' FRAG_A.cif`
    echo "A.inp generated in FRAG_A(top):${FA}; charge and multi: $CHARGEA $MULTIA"
    cd ..

    cd FRAG_B
    gen_minp FRAG_B.cif B.inp $CHARGEB $MULTIB
    FB=`awk '{ if(NF==4 && NR > 16) count[$1]++ } END { for (c in count) printf "%s%s",c,count[c] }' FRAG_B.cif`
    echo "B.inp generated in FRAG_B(base):${FB}; charge and multi: $CHARGEB $MULTIB"
    cd ..

    cd FRAG_AB
    gen_minp POSCAR.cif AB.inp $CHARGEAB $MULTIAB
    echo "AB.inp generated in FRAG_AB; charge and multi: $CHARGEAB $MULTIAB"
    cd ..
  fi
if [[ "$ISDIR" == "1" ]];then
    cd ..
fi
shift 1 
done
